# uamcf/core/pipeline.py
import os
import json
import time
from typing import Dict, List, Any, Tuple, Optional, Union

from ..utils.logger import get_logger
from ..adapters import get_adapter
from ..methods import get_method


class PipelineController:
    """压缩流水线控制器，管理压缩过程的执行"""

    def __init__(self):
        self.logger = get_logger("PipelineController")
        self.config = {}

    def configure(self, config: Dict) -> None:
        """
        配置压缩流水线

        Args:
            config: 配置字典
        """
        self.config = config
        self.logger.info("Pipeline controller configured")

    def execute(self, model: Any, compression_plan: List[Dict],
                adapter_name: str = "pytorch", eval_data: Any = None) -> Tuple[Any, Dict]:
        """
        执行压缩计划

        Args:
            model: 要压缩的模型
            compression_plan: 压缩计划
            adapter_name: 使用的适配器名称
            eval_data: 评估数据

        Returns:
            压缩后的模型和性能统计信息
        """
        self.logger.info(f"Executing compression plan with {len(compression_plan)} steps")

        # 获取适配器
        adapter = get_adapter(adapter_name)
        if adapter is None:
            self.logger.error(f"Unsupported adapter: {adapter_name}")
            return model, {}

        # 记录原始模型信息用于比较
        original_stats = self._measure_model_stats(model, adapter, eval_data)

        # 创建模型副本用于压缩
        compressed_model = adapter.clone_model(model)

        # 记录压缩步骤的统计信息
        step_stats = []

        # 按优先级排序执行压缩步骤
        priority_map = {"high": 3, "medium": 2, "low": 1}
        sorted_plan = sorted(
            compression_plan,
            key=lambda x: priority_map.get(x.get("priority", "low"), 0),
            reverse=True
        )

        # 执行每个压缩步骤
        for i, step in enumerate(sorted_plan):
            step_name = f"Step {i + 1}: {step['method']} on {step['layer']}"
            self.logger.info(f"Executing {step_name}")

            # 获取压缩方法
            method_cls = get_method(step["method"])
            if method_cls is None:
                self.logger.warning(f"Unsupported compression method: {step['method']}")
                continue

            # 创建方法实例
            method_instance = method_cls(step.get("config", {}))

            try:
                # 获取目标层
                target_layer = adapter.get_layer(compressed_model, step["layer"])

                if target_layer is None:
                    self.logger.warning(f"Layer not found: {step['layer']}")
                    continue

                # 记录压缩开始时间
                start_time = time.time()

                # 应用压缩方法
                compressed_layer = method_instance.compress(target_layer)

                # 替换原层
                adapter.replace_layer(compressed_model, step["layer"], compressed_layer)

                # 记录压缩耗时
                elapsed_time = time.time() - start_time

                # 收集步骤统计信息
                step_stat = {
                    "step": i + 1,
                    "layer": step["layer"],
                    "method": step["method"],
                    "config": step.get("config", {}),
                    "time_seconds": elapsed_time
                }

                # 如果提供了评估数据，评估这一步的影响
                if eval_data:
                    step_metrics = self._evaluate_step(compressed_model, adapter, eval_data)
                    step_stat.update(step_metrics)

                step_stats.append(step_stat)
                self.logger.info(f"Completed {step_name} in {elapsed_time:.2f} seconds")

            except Exception as e:
                self.logger.error(f"Error executing {step_name}: {str(e)}")
                continue

        # 如果配置中启用了微调
        if self.config.get("tuning", {}).get("enabled", False) and eval_data:
            self.logger.info("Fine-tuning compressed model")
            compressed_model = self._fine_tune(compressed_model, adapter, eval_data)

        # 测量最终压缩模型的统计信息
        final_stats = self._measure_model_stats(compressed_model, adapter, eval_data)

        # 计算相对于原始模型的改进
        improvement_stats = self._calculate_improvements(original_stats, final_stats)

        # 整合所有统计信息
        stats = {
            "original": original_stats,
            "compressed": final_stats,
            "improvements": improvement_stats,
            "steps": step_stats
        }

        self.logger.info("Compression execution completed")
        return compressed_model, stats

    def _measure_model_stats(self, model: Any, adapter: Any, eval_data: Any = None) -> Dict:
        """测量模型的性能统计信息"""
        stats = {}

        # 模型大小
        stats["size_bytes"] = adapter.get_model_size(model)
        stats["param_count"] = adapter.get_parameter_count(model)

        # 如果有评估数据，测量精度和延迟
        if eval_data:
            try:
                # 评估精度
                accuracy = adapter.evaluate_accuracy(model, eval_data)
                stats["accuracy"] = accuracy

                # 测量延迟（平均推理时间）
                latency = adapter.measure_inference_time(model, eval_data)
                stats["latency_ms"] = latency

                # 测量内存使用
                memory_usage = adapter.measure_memory_usage(model, eval_data)
                stats["memory_mb"] = memory_usage
            except Exception as e:
                self.logger.warning(f"Error measuring model performance: {str(e)}")

        return stats

    def _evaluate_step(self, model: Any, adapter: Any, eval_data: Any) -> Dict:
        """评估单个压缩步骤的影响"""
        metrics = {}

        try:
            # 评估精度
            accuracy = adapter.evaluate_accuracy(model, eval_data)
            metrics["accuracy"] = accuracy

            # 测量延迟
            latency = adapter.measure_inference_time(model, eval_data)
            metrics["latency_ms"] = latency
        except Exception as e:
            self.logger.warning(f"Error evaluating step: {str(e)}")

        return metrics

    def _calculate_improvements(self, original_stats: Dict, final_stats: Dict) -> Dict:
        """计算压缩改进的统计数据"""
        improvements = {}

        # 大小减少
        if "size_bytes" in original_stats and "size_bytes" in final_stats and original_stats["size_bytes"] > 0:
            size_ratio = final_stats["size_bytes"] / original_stats["size_bytes"]
            improvements["size_reduction"] = 1.0 / size_ratio

        # 参数减少
        if "param_count" in original_stats and "param_count" in final_stats and original_stats["param_count"] > 0:
            param_ratio = final_stats["param_count"] / original_stats["param_count"]
            improvements["param_reduction"] = 1.0 / param_ratio

        # 精度变化
        if "accuracy" in original_stats and "accuracy" in final_stats and original_stats["accuracy"] > 0:
            accuracy_change = final_stats["accuracy"] - original_stats["accuracy"]
            improvements["accuracy_change"] = accuracy_change

        # 延迟改进
        if "latency_ms" in original_stats and "latency_ms" in final_stats and original_stats["latency_ms"] > 0:
            latency_ratio = final_stats["latency_ms"] / original_stats["latency_ms"]
            improvements["latency_improvement"] = 1.0 / latency_ratio

        # 内存使用改进
        if "memory_mb" in original_stats and "memory_mb" in final_stats and original_stats["memory_mb"] > 0:
            memory_ratio = final_stats["memory_mb"] / original_stats["memory_mb"]
            improvements["memory_improvement"] = 1.0 / memory_ratio

        return improvements

    def _fine_tune(self, model: Any, adapter: Any, eval_data: Any) -> Any:
        """微调压缩后的模型以恢复精度"""
        self.logger.info("Starting fine-tuning process")

        iterations = self.config.get("tuning", {}).get("iterations", 5)
        early_stopping = self.config.get("tuning", {}).get("early_stopping", True)

        try:
            tuned_model = adapter.fine_tune(
                model,
                eval_data,
                iterations=iterations,
                early_stopping=early_stopping
            )
            self.logger.info("Fine-tuning completed")
            return tuned_model
        except Exception as e:
            self.logger.error(f"Fine-tuning failed: {str(e)}")
            return model

    def export_model(self, model: Any, output_path: str, format: str = "native",
                     adapter_name: str = "pytorch") -> str:
        """
        导出压缩后的模型

        Args:
            model: 要导出的模型
            output_path: 输出路径
            format: 导出格式
            adapter_name: 使用的适配器名称

        Returns:
            导出文件的路径
        """
        adapter = get_adapter(adapter_name)
        if adapter is None:
            self.logger.error(f"Unsupported adapter: {adapter_name}")
            return ""

        try:
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)

            # 导出模型
            export_path = adapter.export_model(model, output_path, format)

            # 导出压缩元数据
            metadata_path = os.path.splitext(output_path)[0] + ".meta.json"
            with open(metadata_path, 'w') as f:
                metadata = {
                    "framework": adapter_name,
                    "format": format,
                    "original_file": output_path,
                    "compression_info": {
                        "version": self.config.get("version", "1.0.0"),
                        "timestamp": time.time(),
                        "stats": self._measure_model_stats(model, adapter)
                    }
                }
                json.dump(metadata, f, indent=2)

            return export_path
        except Exception as e:
            self.logger.error(f"Failed to export model: {str(e)}")
            return ""

    def get_supported_adapters(self) -> Dict[str, str]:
        """获取支持的适配器列表"""
        supported_adapters = {}

        try:
            from ..adapters import ADAPTER_REGISTRY
            for name, adapter_cls in ADAPTER_REGISTRY.items():
                supported_adapters[name] = adapter_cls.__doc__ or f"{name} adapter"
        except Exception:
            pass

        return supported_adapters